import numpy as np
import faiss


# 构建索引，这里我们选用暴力检索的方法FlatL2为例，L2代表构建的index采用的相似度度量方法为L2范数，即欧氏距离
def create_index(datas_embedding, index_path):
    index = faiss.IndexFlatL2(datas_embedding.shape[1])  # 这里必须传入一个向量的维度，创建一个空的索引
    index.add(datas_embedding)  # 把向量数据加入索引
    faiss.write_index(index, index_path)  # 保存索引到文件
    return index

def faiss_index_load(index_path):
    index = faiss.read_index(index_path)
    return index


def get_similar_texts(indices, texts):
    similar_texts = [texts[idx] for idx in indices[0]]
    return similar_texts

def text_to_embedding(txt_path):
    arrays = []

    with open(txt_path, "r", encoding="utf-8") as doc:
        text = doc.read()
        chunks = text.split("|\n")

    for chunk in chunks:
        args = []
        paras = chunk.split("\n")
        for para in paras:
            if para != "":
                arg = float(para)
                if arg != "":
                    args.append(arg)

        if args:
            numpy_array = np.array(args, dtype=np.float64)
            arrays.append(numpy_array)

    numpy_2d_array = np.stack(arrays, axis=0)

    return numpy_2d_array


if __name__ == '__main__':
    # index地址
    index_path = "index/RAG.index"
    header_index_path = "index/header/header_RAG.index"

    # 向量化文本路径
    txt_path = "embedded_files/embedded_text.txt"

    header_path = "embedded_files/embedded_headers.txt"

    # 文本的向量化
    datas_embedding = text_to_embedding(txt_path)
    header_datas_embedding = text_to_embedding(header_path)

    # 创建并保存 FAISS 索引
    create_index(datas_embedding, index_path)
    create_index(header_datas_embedding, header_index_path)

    print(f"向量数据库已成功保存")